
import json
import os
import sqlite3
import openai

import requests

openai.api_key = os.environ["OPENAI_API_KEY_ZHIHU"]
openai.api_base= os.environ["OPENAI_API_BASE_ZHIHU"]


# Example dummy function hard coded to return the same weather
# In production, this could be your backend API or an external API
amap_key = "6d672e6194caa3b639fccf2caf06c342"
cursor = None

table1 = """CREATE TABLE customers (
    id INT PRIMARY KEY NOT NULL, -- 主键，不允许为空
    customer_name VARCHAR(255) NOT NULL, -- 客户名，不允许为空
    email VARCHAR(255) UNIQUE, -- 邮箱，唯一
    register_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP -- 注册时间，默认为当前时间
);"""

table2 = """CREATE TABLE products (
    id INT PRIMARY KEY NOT NULL, -- 主键，不允许为空
    product_name VARCHAR(255) NOT NULL, -- 产品名称，不允许为空
    price DECIMAL(10,2) NOT NULL -- 价格，不允许为空
);"""

table3 = """
CREATE TABLE orders (
    id INT PRIMARY KEY NOT NULL, -- 主键，不允许为空
    customer_id INT NOT NULL, -- 客户ID，不允许为空
    product_id INT NOT NULL, -- 产品ID，不允许为空
    price DECIMAL(10,2) NOT NULL, -- 价格，不允许为空
    status INT NOT NULL, -- 订单状态，整数类型，不允许为空。0代表待支付，1代表已支付，2代表已退款
    create_time TIMESTAMP DEFAULT CURRENT_TIMESTAMP, -- 创建时间，默认为当前时间
    pay_time TIMESTAMP -- 支付时间，可以为空
);
"""
database_schema_string = f"""
{table1}
{table2}
{table3}
"""
tools = [{
    # 摘自 OpenAI 官方示例 https://github.com/openai/openai-cookbook/blob/main/examples/How_to_call_functions_with_chat_models.ipynb
    "type": "function",
    "function": {
        "name": "ask_database",
        "description": "Use this function to answer user questions about business. \
                    Output should be a fully formed SQL query.",
        "parameters": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": f"""
                    SQL query extracting info to answer the user's question.
                    SQL should be written using this database schema:
                    {database_schema_string}
                    The query should be returned in plain text, not in JSON.
                    The query should only contain grammars supported by SQLite.
                    """,
                }
            },
            "required": ["query"],
        }
    }
}]

def getChatResponse(messages,model="gpt-3.5-turbo"):
    # Step 1: send the conversation and available functions to the model
    response = openai.ChatCompletion.create(
        model=model,
        messages=messages,
        temperature=0,
        seed=1024,  # 随机种子保持不变，temperature 和 prompt 不变的情况下，输出就会不变
        tool_choice="auto",  # 默认值，由 GPT 自主决定返回 function call 还是返回文字回复。也可以强制要求必须调用指定的函数，详见官方文档
        tools=tools,
    )
    return response

def run_conversation(userInput,model="gpt-3.5-turbo"):
    messages = [{"role": "system", "content": "你是一个数据分析师，基于数据库的数据回答问题"},
                {"role": "system", "content": userInput}]

    response = getChatResponse(messages,model)
    return response
    # Step 2: check if the model wanted to call a function


result_response = run_conversation("统计每月每件商品的销售额")
sql = json.loads(result_response.choices[0].message.tool_calls[0].function.arguments)["query"]
print(sql)